In [3]:
def explore_prior(N, T, lam_list, mu_list, I0_list):
fig, axes = plt.subplots(len(lam_list), len(mu_list), figsize=(15, 15), sharex=True, sharey=True)
for i, lam in enumerate(lam_list):
for j, mu in enumerate(mu_list):
for I0 in I0_list:
_, states = simulate_sir_simple(lam, mu, I0, T, N)
S,I,R = zip(*states)
S, I, R = np.array(S), np.array(I), np.array(R)
ax=axes[i, j]
ax.plot(range(T), S, label=f'S(t), I0={I0}')
ax.plot(range(T), I, label=f'I(t), I0={I0}')
ax.plot(range(T), R, label=f'R(t), I0={I0}')
ax.set_title(f'λ={lam}, μ={mu}')
ax.set_xlabel('days')
ax.set_ylabel('population')
handles, labels = ax.get_legend_handles_labels()
fig.legend(handles, labels, loc='lower center', bbox_to_anchor=(0.5, -0.05), ncol=len(I0_list))
plt.tight_layout()
plt.show()
In [4]:
N = 1000
T = 100
# hyperparameters
lam_list = np.round(np.linspace(0.1, 0.9, 9),1)
mu_list = np.round(np.linspace(0.1, 0.9, 9),1)
I0_list = [1, 50, 100, 500]
explore_prior(N, T, lam_list, mu_list, I0_list)
With range
- lambda = [0.2, 0.9]
- mu = [0.1, 0.7]
- I_0 = [1, 100]
the most diverse dynamics are observed.
1.3¶
In [5]:
def generate_sir_data(num_samples, T, N):
data = []
target = []
for _ in range(num_samples):
lam = np.random.uniform(0.2, 0.9)
mu = np.random.uniform(0.1, 0.7)
I0 = np.random.randint(1,100)
X, _ = simulate_sir_simple(lam, mu, I0, T, N)
data.append(X.flatten())
target.append([lam, mu, I0])
return np.array(data), np.array(target)
class SIRDataset(Dataset):
def __init__(self, X, Y):
self.X = torch.tensor(X, dtype=torch.float32)
self.Y = torch.tensor(Y, dtype=torch.float32)
def __len__(self):
return len(self.X)
def __getitem__(self, i):
return self.X[i], self.Y[i]
class SIRModel(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(SIRModel, self).__init__()
self.network = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, output_dim)
)
def forward(self, x):
return self.network(x)
In [6]:
def train_model(model, criterion, optimizer, data_loader, epochs):
model.train()
losses = []
for epoch in range(epochs):
epoch_loss = 0.0
for x_batch, y_batch in data_loader:
optimizer.zero_grad()
output = model(x_batch)
loss = criterion(output, y_batch)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
losses.append(epoch_loss/len(data_loader))
if epoch % 5 == 0:
print(f'Epoch {epoch + 1}/{epochs}, Loss: {losses[-1]:.4f}')
return losses
def test_model(model, dataloader):
model.eval()
preds = []
targets = []
with torch.no_grad():
for x_batch, y_batch in dataloader:
outputs = model(x_batch)
preds.append(np.array(outputs))
targets.append(np.array(y_batch))
return np.vstack(preds), np.vstack(targets)
In [7]:
# hyperparameters
N = 1000
T = 100
num_samples = 5000
test_samples = 1000
input_size = T*2 # flattened SIR
hidden_size = 128
output_size = 3 # [lam, mu, I0]
batch_size = 64
lr = 0.001
epochs = 50
# generate data
x_train, y_train = generate_sir_data(num_samples, T, N)
x_test, y_test = generate_sir_data(num_samples=test_samples, T=T, N=N)
# preprocess data
train_dataset = SIRDataset(x_train, y_train)
test_dataset = SIRDataset(x_test, y_test)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# initialize model
model = SIRModel(input_size, hidden_size, output_size)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
# train model
print('Training mode ...')
train_loss = train_model(model, criterion, optimizer, train_loader, epochs)
# visualize training loss
plt.figure(figsize=(8,6))
sns.lineplot(x=np.arange(1, epochs+1), y=train_loss)
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Reconstruction Error')
plt.show()
#test model
print('Test mode ...')
preds, targets = test_model(model, test_loader)
# visualize predictions vs targets
cols = ['λ', 'μ', 'I_0']
pred_df = pd.DataFrame(preds, columns=cols)
target_df = pd.DataFrame(targets, columns=cols)
fig, axes = plt.subplots(1, 3, figsize=(15,5))
for i, param in enumerate(cols):
sns.scatterplot(x=target_df[param], y=pred_df[param], alpha=0.7, ax=axes[i])
axes[i].plot([target_df[param].min(), target_df[param].max()],
[target_df[param].min(), target_df[param].max()],
color='red', linestyle='--')
axes[i].set_xlabel(f'True {param}')
axes[i].set_ylabel(f'Predicted {param}')
axes[i].set_title(f'Target vs Predicted for {param}')
plt.tight_layout()
plt.show()
Training mode ... Epoch 1/50, Loss: 241.9659 Epoch 6/50, Loss: 8.5415 Epoch 11/50, Loss: 2.2789 Epoch 16/50, Loss: 0.5541 Epoch 21/50, Loss: 0.2494 Epoch 26/50, Loss: 0.4930 Epoch 31/50, Loss: 0.3878 Epoch 36/50, Loss: 0.2304 Epoch 41/50, Loss: 0.5943 Epoch 46/50, Loss: 0.5427
Test mode ...
/var/folders/gl/m711nqcx42d81f7zr656_v600000gn/T/ipykernel_47538/4015364856.py:30: DeprecationWarning: __array__ implementation doesn't accept a copy keyword, so passing copy=False failed. __array__ must implement 'dtype' and 'copy' keyword arguments. preds.append(np.array(outputs)) /var/folders/gl/m711nqcx42d81f7zr656_v600000gn/T/ipykernel_47538/4015364856.py:31: DeprecationWarning: __array__ implementation doesn't accept a copy keyword, so passing copy=False failed. __array__ must implement 'dtype' and 'copy' keyword arguments. targets.append(np.array(y_batch))
1.4¶
In [8]:
# Conditional realNVP components
class ConditionalCouplingLayer(nn.Module):
def __init__(self, input_size, condition_size, hidden_size):
super(ConditionalCouplingLayer, self).__init__()
self.split1 = input_size // 2
self.split2 = input_size - self.split1
self.subnet = nn.Sequential(
nn.Linear(self.split1 + condition_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, 2 * self.split2),
)
def forward(self, x, cond):
x1, x2 = x[..., :self.split1], x[..., self.split1:]
params = self.subnet(torch.cat([x1, cond], dim=-1))
s, t = params[..., :self.split2], params[..., self.split2:]
s = torch.tanh(s)
log_jacob_det = torch.sum(s, dim=-1)
x2 = torch.exp(s) * x2 + t
return torch.cat([x1, x2], dim=-1), log_jacob_det
def reverse(self, x, cond):
x1, x2 = x[..., :self.split1], x[..., self.split1:]
params = self.subnet(torch.cat([x1, cond], dim=-1))
s, t = params[..., :self.split2], params[..., self.split2:]
s = torch.tanh(s)
x2 = (x2 - t) * torch.exp(-s)
return torch.cat([x1, x2], dim=-1)
class ConditionalrealNVP(nn.Module):
def __init__(self, input_size, condition_size, hidden_size, n_layers):
super(ConditionalrealNVP, self).__init__()
self.layers = nn.ModuleList([
ConditionalCouplingLayer(input_size, condition_size, hidden_size) for _ in range(n_layers)
])
def forward(self, x, context):
log_det_jacobian = 0
for layer in self.layers:
x, ldj = layer(x, context)
log_det_jacobian += ldj
return x, log_det_jacobian
def reverse(self, z, context):
for layer in reversed(self.layers):
z = layer.reverse(z, context)
return z
In [9]:
def train_realNVP(feature_extractor, realNVP, optimizer, data_loader, epochs):
feature_extractor.train()
realNVP.train()
losses = []
for epoch in range(epochs):
epoch_loss = 0.0
for x_batch, y_batch in data_loader:
optimizer.zero_grad()
context = feature_extractor(x_batch)
z, log_det_jacobian = realNVP(y_batch, context)
nll = 0.5 * torch.sum(z**2, dim=1) - log_det_jacobian
loss = torch.mean(nll)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
losses.append(epoch_loss/len(data_loader))
if epoch % 2 == 0:
print(f'Epoch {epoch + 1}/{epochs}, Loss: {losses[-1]:.4f}')
return np.array(losses)
def test_realNVP(feature_extractor, realNVP, data_loader):
feature_extractor.eval()
realNVP.eval()
preds = []
targets = []
with torch.no_grad():
for x_batch, y_batch in data_loader:
context = feature_extractor(x_batch)
outputs = realNVP.reverse(torch.randn_like(y_batch), context)
preds.append(np.array(outputs))
targets.append(np.array(y_batch))
return np.vstack(preds), np.vstack(targets)
In [10]:
# Hyperparameters
T = 100
N = 1000
num_samples = 5000
test_samples = 1000
input_size = T*2 # Flattened (dS, dR)
hidden_size = 64
batch_size = 64
lr = 0.001
epochs = 50
n_layers=2
# generate data
x_train, y_train = generate_sir_data(num_samples, T, N)
x_test, y_test = generate_sir_data(num_samples=test_samples, T=T, N=N)
# preprocess data
train_dataset = SIRDataset(x_train, y_train)
test_dataset = SIRDataset(x_test, y_test)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# Initialize feature extractor and realNVP
feature_extractor = SIRModel(input_size, hidden_size, hidden_size)
realNVP = ConditionalrealNVP(input_size=3, condition_size=hidden_size, hidden_size=hidden_size, n_layers=n_layers)
# Define optimizer
params = list(feature_extractor.parameters()) + list(realNVP.parameters())
optimizer = optim.Adam(params, lr=lr)
# train model
print('Training mode ...')
train_loss = train_realNVP(feature_extractor, realNVP, optimizer, train_loader, epochs)
# visualize training loss
plt.figure(figsize=(8,6))
sns.lineplot(x=np.arange(1, epochs+1), y=train_loss)
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Reconstruction Error')
plt.show()
# test model
print('Test mode ...')
preds, targets = test_realNVP(feature_extractor, realNVP, test_loader)
# visualize predictions vs targets
cols = ['λ', 'μ', 'I_0']
pred_df = pd.DataFrame(preds, columns=cols)
target_df = pd.DataFrame(targets, columns=cols)
fig, axes = plt.subplots(1, 3, figsize=(15,5))
for i, param in enumerate(cols):
sns.scatterplot(x=target_df[param], y=pred_df[param], alpha=0.7, ax=axes[i])
axes[i].plot([target_df[param].min(), target_df[param].max()],
[target_df[param].min(), target_df[param].max()],
color='red', linestyle='--')
axes[i].set_xlabel(f'True {param}')
axes[i].set_ylabel(f'Predicted {param}')
axes[i].set_title(f'Posterior for {param}')
plt.tight_layout()
plt.show()
Training mode ... Epoch 1/50, Loss: 63.7104 Epoch 3/50, Loss: 1.4205 Epoch 5/50, Loss: 1.0354 Epoch 7/50, Loss: 0.6610 Epoch 9/50, Loss: 0.4803 Epoch 11/50, Loss: -0.0064 Epoch 13/50, Loss: 0.0158 Epoch 15/50, Loss: -0.2190 Epoch 17/50, Loss: -0.0612 Epoch 19/50, Loss: 0.0826 Epoch 21/50, Loss: -0.4647 Epoch 23/50, Loss: -0.0404 Epoch 25/50, Loss: 0.2591 Epoch 27/50, Loss: -0.4760 Epoch 29/50, Loss: -0.1908 Epoch 31/50, Loss: -0.2632 Epoch 33/50, Loss: 0.4863 Epoch 35/50, Loss: 0.2263 Epoch 37/50, Loss: -0.0194 Epoch 39/50, Loss: -0.1643 Epoch 41/50, Loss: -0.1991 Epoch 43/50, Loss: -0.1115 Epoch 45/50, Loss: 0.2068 Epoch 47/50, Loss: -0.3333 Epoch 49/50, Loss: -0.3454
Test mode ...
/var/folders/gl/m711nqcx42d81f7zr656_v600000gn/T/ipykernel_47538/443406207.py:37: DeprecationWarning: __array__ implementation doesn't accept a copy keyword, so passing copy=False failed. __array__ must implement 'dtype' and 'copy' keyword arguments. preds.append(np.array(outputs)) /var/folders/gl/m711nqcx42d81f7zr656_v600000gn/T/ipykernel_47538/443406207.py:38: DeprecationWarning: __array__ implementation doesn't accept a copy keyword, so passing copy=False failed. __array__ must implement 'dtype' and 'copy' keyword arguments. targets.append(np.array(y_batch))
1.5¶
In [11]:
def generate_posterior_samples(feature_extractor, realNVP, test_loader, num_preds, input_size):
feature_extractor.eval()
realNVP.eval()
all_samples = []
all_targets = []
with torch.no_grad():
for x, y in test_loader:
x = x.to(next(feature_extractor.parameters()).device)
context = feature_extractor(x)
samples = []
for _ in range(num_preds):
z = torch.randn((x.size(0), input_size), device=context.device) # Sampling from prior
sample = realNVP.reverse(z, context)
samples.append(sample.cpu().numpy())
all_samples.append(np.stack(samples, axis=0))
all_targets.append(y.numpy())
preds = np.concatenate(all_samples, axis=1) # Shape: (M, len(test_data), num_params)
targets = np.concatenate(all_targets, axis=0) # Shape: (len(test_data), num_params)
return preds, targets
def compute_energy(y_samples, targets):
targets = targets[None, :, :] # Broadcast for M samples
energy = -0.5 * np.sum((y_samples - targets) ** 2, axis=-1) # Gaussian assumption
return energy
def plot_posterior_distributions(samples, params):
fig, axes = plt.subplots(2, len(params), figsize=(15, 5))
for i, param in enumerate(params):
param_samples = samples[:, :, i].flatten()
axes[0,i].hist(param_samples, bins=30, density=True, alpha=0.5, label='Posterior Samples')
axes[0,i].set_title(f'Posterior Histogram - {param}')
axes[0,i].set_xlabel('Value')
axes[0,i].set_ylabel('Density')
axes[0,i].legend()
# Empirical CDF
ecdf = np.sort(param_samples)
axes[1,i].plot(ecdf, np.linspace(0, 1, len(ecdf)), label='Empirical CDF', lw=2)
axes[1,i].set_title(f'Empirical CDF - {param}')
axes[1,i].set_xlabel('Value')
axes[1,i].set_ylabel('Cumulative Probability')
axes[1,i].legend()
plt.tight_layout()
plt.show()
def plot_joint_distribution(samples, params):
joint_samples = samples.reshape(-1, samples.shape[-1]) # Shape: (M * len(X), num_params)
df = pd.DataFrame(joint_samples, columns=params)
sns.pairplot(df, diag_kind='kde', plot_kws={'s': 1, 'alpha':0.5}, corner=True)
plt.suptitle('Joint Posterior Distribution', y=1)
plt.show()
In [12]:
num_preds = 100
input_size = 3
params = ['λ', 'μ', 'I_0']
preds, targets = generate_posterior_samples(feature_extractor, realNVP, test_loader, num_preds, input_size)
energy_scores = compute_energy(preds, targets)
# displaying summarized energy scores
print(f'Mean Energy Score: {np.mean(energy_scores):.4f}')
print(f'Standard Deviation of Energy Scores: {np.std(energy_scores):.4f}')
# marginal distributions
plot_posterior_distributions(preds, params)
# joint distributions
plot_joint_distribution(preds, params)
Mean Energy Score: -8.0579 Standard Deviation of Energy Scores: 11.5818
1.6¶
In [13]:
def plot_confidence_intervals(resimulated_outcomes, true_simulated_outcomes, coverage):
lower_bound = np.percentile(resimulated_outcomes, 2.5, axis=0)
upper_bound = np.percentile(resimulated_outcomes, 97.5, axis=0)
params = ['dS', 'dR']
fig, axes = plt.subplots(1, 2, figsize=(15, 5), sharex=True, sharey=True)
for i, param in enumerate(params):
axes[i].fill_between(range(len(lower_bound)), lower_bound[:, i], upper_bound[:, i], color='skyblue', alpha=0.5, label='95% CI')
axes[i].plot(true_simulated_outcomes[:, i], 'o', color='orange', label='True Outcomes')
axes[i].set_xlabel('samples')
axes[i].set_ylabel(f'{param}')
axes[i].legend()
plt.suptitle(f'Posterior Predictive Calibration with coverage = {coverage}')
plt.tight_layout()
plt.show()
In [14]:
def posterior_predictive_calibration(posterior_samples, simulate_sir_simple, true_outcomes, T, N):
M = posterior_samples.shape[0]
resimulated_outcomes = []
# simulate outcomes for predicted samples
for k in range(M):
current_outcomes = []
for lam, mu, I0 in posterior_samples[k]:
X, _ = simulate_sir_simple(lam, mu, I0, T, N)
current_outcomes.append(np.sum(X, axis=0))
resimulated_outcomes.append(np.stack(current_outcomes, axis=0))
resimulated_outcomes = np.stack(resimulated_outcomes, axis=0)
# simulate outcomes for true samples
true_simulated_outcomes = []
for lam, mu, I0 in true_outcomes:
X, _ = simulate_sir_simple(lam, mu, I0, T, N)
true_simulated_outcomes.append(np.sum(X, axis=0))
true_simulated_outcomes = np.stack(true_simulated_outcomes, axis=0)
# confidence intervals from predicted simulations
lower_bound = np.percentile(resimulated_outcomes, 2.5, axis=0)
upper_bound = np.percentile(resimulated_outcomes, 97.5, axis=0)
# Check if true simulated outcomes are within confidence intervals
true_within_interval = np.logical_and(true_simulated_outcomes >= lower_bound, true_simulated_outcomes <= upper_bound)
# coverage as the proportion of observations within intervals
coverage = np.mean(np.all(true_within_interval, axis=1))
plot_confidence_intervals(resimulated_outcomes, true_simulated_outcomes, coverage)
return coverage
In [15]:
posterior_predictive_calibration(preds, simulate_sir_simple, y_test, T, N)
/var/folders/gl/m711nqcx42d81f7zr656_v600000gn/T/ipykernel_47538/1062177979.py:12: RuntimeWarning: overflow encountered in scalar multiply dS = -lam * S * I / N /var/folders/gl/m711nqcx42d81f7zr656_v600000gn/T/ipykernel_47538/1062177979.py:13: RuntimeWarning: overflow encountered in scalar multiply dI = lam * S * I / N - (mu * I) /var/folders/gl/m711nqcx42d81f7zr656_v600000gn/T/ipykernel_47538/1062177979.py:21: RuntimeWarning: invalid value encountered in scalar subtract dS = S - S_next /var/folders/gl/m711nqcx42d81f7zr656_v600000gn/T/ipykernel_47538/1062177979.py:22: RuntimeWarning: invalid value encountered in scalar subtract dR = R_next - R /var/folders/gl/m711nqcx42d81f7zr656_v600000gn/T/ipykernel_47538/1062177979.py:13: RuntimeWarning: invalid value encountered in scalar subtract dI = lam * S * I / N - (mu * I) /Users/sandrinjoy/Desktop/Dacs/gnns-ws-24/.conda/lib/python3.10/site-packages/numpy/lib/_function_base_impl.py:4623: RuntimeWarning: invalid value encountered in subtract subtract(b, diff_b_a * (1 - t), out=lerp_interpolation, where=t >= 0.5,
Out[15]:
np.float64(0.32)
Task 2¶
2.1¶
In [16]:
import numpy as np
def simulate_sir_noisy(lam, mu, I0, L, rho, sigma2, T, N):
X, C = simulate_sir_simple(lam, mu, I0, T, N)
# multiplicative Gaussian noise
ep_S = np.random.normal(rho, np.sqrt(sigma2), T)
ep_I = np.random.normal(rho, np.sqrt(sigma2), T)
ep_R = np.random.normal(rho, np.sqrt(sigma2), T)
dS, dR = X[:, 0]*ep_S, X[:, 1]*ep_R
S, I, R = C[:,0]*ep_S, C[:,1]*ep_I, C[:,2]*ep_R
X = np.stack([dS, dR], axis=1)
C = np.stack([S, I, R], axis=1)
X_shifted = np.roll(X, int(L), axis=0)
C_shifted = np.roll(C, int(L), axis=0)
return X_shifted, C_shifted
In [17]:
def explore_noisy_prior(lam, mu, I0, N, T, L_list, rho_list, sigma_list):
fig, axes = plt.subplots(len(rho_list), len(sigma_list), figsize=(15, 15), sharex=True, sharey=True)
for i, rho in enumerate(rho_list):
for j, sigma2 in enumerate(sigma_list):
for L in L_list:
sim, _ = simulate_sir_noisy(lam, mu, I0, L, rho, sigma2, T, N)
S, R = zip(*sim)
S, R = np.array(S), np.array(R)
ax=axes[i, j]
ax.plot(range(T), S, label=f'S(t), L={L}', linewidth=0.5)
ax.plot(range(T), R, label=f'R(t), L={L}', linewidth=0.5)
ax.set_title(f'ρ={rho}, σ2={sigma2}')
ax.set_xlabel('days')
ax.set_ylabel('value')
handles, labels = ax.get_legend_handles_labels()
fig.legend(handles, labels, loc='lower center', bbox_to_anchor=(0.5, -0.05), ncol=len(L_list))
plt.tight_layout()
plt.show()
In [18]:
lam=0.4
mu=0.2
I0=1
N = 1000
T = 100
# hyperparameters
rho_list = np.round(np.linspace(0.1, 0.9, 9), 1)
sigma_list = np.round(np.linspace(0.1, 0.9, 9), 1)
L_list = np.linspace(0, 10, 2)
In [19]:
explore_noisy_prior(lam, mu, I0, N, T, L_list, rho_list, sigma_list)
The picked ranges are
- rho = [0.1, 0.9]
- sigma = [0.3, 0.9]
- L = [1, 10]
the most diverse dynamics are observed.
In [20]:
def generate_noisy_sir_data(num_samples, lam, mu, I0, T, N):
data = []
target = []
for _ in range(num_samples):
rho = np.random.uniform(0.1, 0.9)
sigma2 = np.random.uniform(0.3, 0.9)
L = np.random.randint(1,10)
X, _ = simulate_sir_noisy(lam, mu, I0, L, rho, sigma2, T, N)
data.append(X.flatten())
target.append([lam, mu, I0, rho, sigma2, L])
return np.array(data), np.array(target)
In [21]:
# hyperparameters
lam=0.4
mu=0.2
I0=1
N = 1000
T = 100
num_samples = 5000
test_samples = 1000
input_size = T*2 # flattened SIR
hidden_size = 128
output_size = 6 # [lam, mu, I0, rho, sigma2, L]
batch_size = 64
lr = 0.001
epochs = 50
# generate data
x_train, y_train = generate_noisy_sir_data(num_samples, lam, mu, I0, T, N)
x_test, y_test = generate_noisy_sir_data(num_samples=test_samples, lam=lam, mu=mu, I0=I0, T=T, N=N)
# preprocess data
train_dataset = SIRDataset(x_train, y_train)
test_dataset = SIRDataset(x_test, y_test)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# initialize model
model = SIRModel(input_size, hidden_size, output_size)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
# train model
print('Training mode ...')
train_loss = train_model(model, criterion, optimizer, train_loader, epochs)
# visualize training loss
plt.figure(figsize=(8,6))
sns.lineplot(x=np.arange(1, epochs+1), y=train_loss)
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Reconstruction Error')
plt.show()
#test model
print('Test mode ...')
preds, targets = test_model(model, test_loader)
# visualize predictions vs targets
cols = ['λ', 'μ', 'I_0', 'ρ', 'σ2', 'L']
pred_df = pd.DataFrame(preds, columns=cols)
target_df = pd.DataFrame(targets, columns=cols)
fig, axes = plt.subplots(1, len(cols), figsize=(15,5))
for i, param in enumerate(cols):
sns.scatterplot(x=target_df[param], y=pred_df[param], alpha=0.7, ax=axes[i])
axes[i].plot([target_df[param].min(), target_df[param].max()],
[target_df[param].min(), target_df[param].max()],
color='red', linestyle='--')
axes[i].set_xlabel(f'True {param}')
axes[i].set_ylabel(f'Predicted {param}')
axes[i].set_title(f'Target vs Predicted for {param}')
plt.tight_layout()
plt.show()
Training mode ... Epoch 1/50, Loss: 1.0963 Epoch 6/50, Loss: 0.1832 Epoch 11/50, Loss: 0.0823 Epoch 16/50, Loss: 0.0445 Epoch 21/50, Loss: 0.0410 Epoch 26/50, Loss: 0.0318 Epoch 31/50, Loss: 0.0170 Epoch 36/50, Loss: 0.0182 Epoch 41/50, Loss: 0.0278 Epoch 46/50, Loss: 0.0166
Test mode ...
/var/folders/gl/m711nqcx42d81f7zr656_v600000gn/T/ipykernel_47538/4015364856.py:30: DeprecationWarning: __array__ implementation doesn't accept a copy keyword, so passing copy=False failed. __array__ must implement 'dtype' and 'copy' keyword arguments. preds.append(np.array(outputs)) /var/folders/gl/m711nqcx42d81f7zr656_v600000gn/T/ipykernel_47538/4015364856.py:31: DeprecationWarning: __array__ implementation doesn't accept a copy keyword, so passing copy=False failed. __array__ must implement 'dtype' and 'copy' keyword arguments. targets.append(np.array(y_batch))
Conditional RealNVP
In [22]:
# Hyperparameters
input_size = T*2 # Flattened (dS, dR)
hidden_size = 128
batch_size = 64
lr = 0.001
epochs = 100
n_layers=2
# Initialize feature extractor and realNVP
feature_extractor = SIRModel(input_size, hidden_size, hidden_size)
realNVP = ConditionalrealNVP(input_size=6, condition_size=hidden_size, hidden_size=hidden_size, n_layers=n_layers)
# Define optimizer
params = list(feature_extractor.parameters()) + list(realNVP.parameters())
optimizer = optim.Adam(params, lr=lr)
# train model
print('Training mode ...')
train_loss = train_realNVP(feature_extractor, realNVP, optimizer, train_loader, epochs)
# visualize training loss
plt.figure(figsize=(8,6))
sns.lineplot(x=np.arange(1, epochs+1), y=train_loss)
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Reconstruction Error')
plt.show()
# test model
print('Test mode ...')
preds, targets = test_realNVP(feature_extractor, realNVP, test_loader)
# visualize predictions vs targets
cols = ['λ', 'μ', 'I_0', 'ρ', 'σ2', 'L']
pred_df = pd.DataFrame(preds, columns=cols)
target_df = pd.DataFrame(targets, columns=cols)
fig, axes = plt.subplots(1, len(cols), figsize=(15,5))
for i, param in enumerate(cols):
sns.scatterplot(x=target_df[param], y=pred_df[param], alpha=0.7, ax=axes[i])
axes[i].plot([target_df[param].min(), target_df[param].max()],
[target_df[param].min(), target_df[param].max()],
color='red', linestyle='--')
axes[i].set_xlabel(f'True {param}')
axes[i].set_ylabel(f'Predicted {param}')
axes[i].set_title(f'Posterior for {param}')
plt.tight_layout()
plt.show()
Training mode ... Epoch 1/100, Loss: -0.0467 Epoch 3/100, Loss: -0.8726 Epoch 5/100, Loss: -2.1709 Epoch 7/100, Loss: -2.5259 Epoch 9/100, Loss: -2.7570 Epoch 11/100, Loss: -2.9374 Epoch 13/100, Loss: -3.0440 Epoch 15/100, Loss: -3.1015 Epoch 17/100, Loss: -3.1921 Epoch 19/100, Loss: -3.2565 Epoch 21/100, Loss: -3.2594 Epoch 23/100, Loss: -3.4279 Epoch 25/100, Loss: -3.5185 Epoch 27/100, Loss: -3.0727 Epoch 29/100, Loss: -3.6016 Epoch 31/100, Loss: -3.6795 Epoch 33/100, Loss: -3.6779 Epoch 35/100, Loss: -3.8543 Epoch 37/100, Loss: -3.8682 Epoch 39/100, Loss: -3.9784 Epoch 41/100, Loss: -3.9048 Epoch 43/100, Loss: -4.0113 Epoch 45/100, Loss: -4.0743 Epoch 47/100, Loss: -4.1614 Epoch 49/100, Loss: -4.1267 Epoch 51/100, Loss: -4.1703 Epoch 53/100, Loss: -4.2428 Epoch 55/100, Loss: -4.1957 Epoch 57/100, Loss: -4.2576 Epoch 59/100, Loss: -4.1049 Epoch 61/100, Loss: -4.1215 Epoch 63/100, Loss: -4.2626 Epoch 65/100, Loss: -4.3307 Epoch 67/100, Loss: -4.3653 Epoch 69/100, Loss: -4.4095 Epoch 71/100, Loss: -4.2955 Epoch 73/100, Loss: -4.1043 Epoch 75/100, Loss: -4.2600 Epoch 77/100, Loss: -4.4253 Epoch 79/100, Loss: -4.2974 Epoch 81/100, Loss: -4.4564 Epoch 83/100, Loss: -4.5650 Epoch 85/100, Loss: -4.3737 Epoch 87/100, Loss: -4.3488 Epoch 89/100, Loss: -4.5145 Epoch 91/100, Loss: -4.5115 Epoch 93/100, Loss: -4.2812 Epoch 95/100, Loss: -4.5263 Epoch 97/100, Loss: -4.5606 Epoch 99/100, Loss: -4.5183
Test mode ...
/var/folders/gl/m711nqcx42d81f7zr656_v600000gn/T/ipykernel_47538/443406207.py:37: DeprecationWarning: __array__ implementation doesn't accept a copy keyword, so passing copy=False failed. __array__ must implement 'dtype' and 'copy' keyword arguments. preds.append(np.array(outputs)) /var/folders/gl/m711nqcx42d81f7zr656_v600000gn/T/ipykernel_47538/443406207.py:38: DeprecationWarning: __array__ implementation doesn't accept a copy keyword, so passing copy=False failed. __array__ must implement 'dtype' and 'copy' keyword arguments. targets.append(np.array(y_batch))
posterior distributions
In [23]:
num_preds = 100
input_size = 6
params = ['λ', 'μ', 'I_0', 'ρ', 'σ2', 'L']
preds, targets = generate_posterior_samples(feature_extractor, realNVP, test_loader, num_preds, input_size)
energy_scores = compute_energy(preds, targets)
# displaying summarized energy scores
print(f'Mean Energy Score: {np.mean(energy_scores):.4f}')
print(f'Standard Deviation of Energy Scores: {np.std(energy_scores):.4f}')
# marginal distributions
plot_posterior_distributions(preds, params)
# joint distributions
plot_joint_distribution(preds, params)
Mean Energy Score: -2.7437 Standard Deviation of Energy Scores: 2.0943
TASK 3¶
In [24]:
def train_model(model, dataloader, criterion, optimizer, num_epochs):
model.train()
loss_history = []
for epoch in range(num_epochs):
epoch_loss = 0.0
for inputs, targets in dataloader:
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
loss_history.append(epoch_loss / len(dataloader))
return loss_history
def evaluate_model(model, dataloader, criterion):
model.eval()
total_loss = 0.0
with torch.no_grad():
for inputs, targets in dataloader:
outputs = model(inputs)
loss = criterion(outputs, targets)
total_loss += loss.item()
return total_loss / len(dataloader)
Simluate Data¶
In [25]:
# Generate test data using simulate_sir_simple()
lam = 0.3
mu = 0.1
I0 = 10
T = 100
N = 1000
simple_data, simple_states = simulate_sir_simple(lam, mu, I0, T, N)
# Generate test data using simulate_sir_noisy()
noise_level = 0.1
sigma2 = noise_level**2
L = 1
noisy_data, noisy_states = simulate_sir_noisy(lam, mu, I0, L, noise_level, sigma2, T, N)
# Display the shapes of the generated data
print("Simple data shape:", simple_data.shape)
print("Noisy data shape:", noisy_data.shape)
# Plot a few samples of the generated data
import matplotlib.pyplot as plt
plt.figure(figsize=(12, 6))
# Plot simple data
plt.subplot(1, 2, 1)
plt.plot(simple_data[:5].T)
plt.title("Simple SIR Data")
plt.xlabel("Timesteps")
plt.ylabel("Values")
# Plot noisy data
plt.subplot(1, 2, 2)
plt.plot(noisy_data[:5].T)
plt.title("Noisy SIR Data")
plt.xlabel("Timesteps")
plt.ylabel("Values")
plt.tight_layout()
plt.show()
Simple data shape: (100, 2) Noisy data shape: (100, 2)
Train Networks with NLL and MMD Loss¶
Modify the training process to include both NLL and MMD loss on the output of the summary network h(X).
In [26]:
# Define the MMD loss function
def mmd_loss(output, target):
# Compute the MMD loss between the output and the target
xx, yy, zz = torch.mm(output, output.t()), torch.mm(target, target.t()), torch.mm(output, target.t())
rx = (xx.diag().unsqueeze(0).expand_as(xx))
ry = (yy.diag().unsqueeze(0).expand_as(yy))
K = torch.exp(-0.5 * (rx.t() + rx - 2*xx))
L = torch.exp(-0.5 * (ry.t() + ry - 2*yy))
P = torch.exp(-0.5 * (rx.t() + ry - 2*zz))
beta = (1./(output.size(0)*output.size(0)))
gamma = (2./(output.size(0)*target.size(0)))
return beta * (K.sum() + L.sum()) - gamma * P.sum()
# Define the neural network model
class SummaryNetwork(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(SummaryNetwork, self).__init__()
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.fc2 = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# Prepare the data
input_dim = simple_data.shape[1]
hidden_dim = 128
output_dim = 10
num_samples = simple_data.shape[0]
# Convert data to PyTorch tensors
simple_data_tensor = torch.tensor(simple_data, dtype=torch.float32)
noisy_data_tensor = torch.tensor(noisy_data, dtype=torch.float32)
# Create DataLoader for simple data
simple_dataset = TensorDataset(simple_data_tensor, torch.zeros(num_samples, output_dim))
simple_dataloader = DataLoader(simple_dataset, batch_size=32, shuffle=True)
# Initialize the model, loss functions, and optimizer
model = SummaryNetwork(input_dim, hidden_dim, output_dim)
nll_loss = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# Train the model with NLL and MMD loss
num_epochs = 50
loss_history = []
for epoch in range(num_epochs):
model.train()
epoch_loss = 0.0
for inputs, _ in simple_dataloader:
optimizer.zero_grad()
outputs = model(inputs)
# Generate standard normal target
target = torch.randn_like(outputs)
# Compute NLL and MMD loss
loss_nll = nll_loss(outputs, target)
loss_mmd = mmd_loss(outputs, target)
loss = loss_nll + loss_mmd
loss.backward()
optimizer.step()
epoch_loss += loss.item()
loss_history.append(epoch_loss / len(simple_dataloader))
print(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss / len(simple_dataloader)}")
# Plot the training loss history
plt.figure(figsize=(8, 4))
plt.plot(loss_history, label='Training Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training Loss History')
plt.legend()
plt.show()
Epoch 1/50, Loss: 13.31803274154663 Epoch 2/50, Loss: 5.530503273010254 Epoch 3/50, Loss: 3.907529890537262 Epoch 4/50, Loss: 2.5789492428302765 Epoch 5/50, Loss: 2.396752715110779 Epoch 6/50, Loss: 2.2364628314971924 Epoch 7/50, Loss: 2.028217315673828 Epoch 8/50, Loss: 1.8717308044433594 Epoch 9/50, Loss: 1.9078839719295502 Epoch 10/50, Loss: 1.8851113319396973 Epoch 11/50, Loss: 1.759248971939087 Epoch 12/50, Loss: 1.764386922121048 Epoch 13/50, Loss: 1.8810383081436157 Epoch 14/50, Loss: 1.772846758365631 Epoch 15/50, Loss: 1.7020862400531769 Epoch 16/50, Loss: 1.7379138469696045 Epoch 17/50, Loss: 1.7030720114707947 Epoch 18/50, Loss: 1.6425380110740662 Epoch 19/50, Loss: 1.6355963051319122 Epoch 20/50, Loss: 1.7186952531337738 Epoch 21/50, Loss: 1.69977268576622 Epoch 22/50, Loss: 1.5730359554290771 Epoch 23/50, Loss: 1.702895849943161 Epoch 24/50, Loss: 1.5550656914710999 Epoch 25/50, Loss: 1.6368052959442139 Epoch 26/50, Loss: 1.7828851342201233 Epoch 27/50, Loss: 1.6247079968452454 Epoch 28/50, Loss: 1.7217609882354736 Epoch 29/50, Loss: 1.645445078611374 Epoch 30/50, Loss: 1.70895054936409 Epoch 31/50, Loss: 1.7866719961166382 Epoch 32/50, Loss: 1.6241732835769653 Epoch 33/50, Loss: 1.7292563021183014 Epoch 34/50, Loss: 1.5987822711467743 Epoch 35/50, Loss: 1.4555167257785797 Epoch 36/50, Loss: 1.584316998720169 Epoch 37/50, Loss: 1.727157711982727 Epoch 38/50, Loss: 1.6102245151996613 Epoch 39/50, Loss: 1.647591382265091 Epoch 40/50, Loss: 1.6078676581382751 Epoch 41/50, Loss: 1.5434271395206451 Epoch 42/50, Loss: 1.5916908979415894 Epoch 43/50, Loss: 1.6677139699459076 Epoch 44/50, Loss: 1.5955264866352081 Epoch 45/50, Loss: 1.652401626110077 Epoch 46/50, Loss: 1.6758142113685608 Epoch 47/50, Loss: 1.5622536838054657 Epoch 48/50, Loss: 1.6523127555847168 Epoch 49/50, Loss: 1.5256880819797516 Epoch 50/50, Loss: 1.5343037843704224
Evaluate Null Hypothesis with Noise-Free Data¶
Use test data generated by simulate_sir_simple() to ensure that the null hypothesis cannot be rejected for noise-free data at common p-values.
In [27]:
# Generate outputs from the trained model
model.eval()
with torch.no_grad():
outputs = model(simple_data_tensor)
# Perform Chi-Square test
chi2_stat, p_value = stats.chisquare(outputs.numpy().flatten())
# Print the Chi-Square test results
print(f"Chi-Square Statistic: {chi2_stat}")
print(f"P-Value: {p_value}")
# Check if the null hypothesis can be rejected at 1% significance level
alpha = 0.01
if p_value < alpha:
print("Reject the null hypothesis: The data does not follow the standard normal distribution.")
else:
print("Fail to reject the null hypothesis: The data follows the standard normal distribution.")
Chi-Square Statistic: 5191.577608035008 P-Value: 0.0 Reject the null hypothesis: The data does not follow the standard normal distribution.
Evaluate Model with Noisy Data¶
Use test data created with simulate_sir_noisy() to determine the conditions under which the model mismatch is detected by the χ2 test.
In [28]:
# Prepare the noisy data
noisy_dataset = TensorDataset(noisy_data_tensor, torch.zeros(num_samples, output_dim))
noisy_dataloader = DataLoader(noisy_dataset, batch_size=32, shuffle=False)
# Generate outputs from the trained model using noisy data
model.eval()
with torch.no_grad():
noisy_outputs = model(noisy_data_tensor)
# Perform Chi-Square test on noisy data
chi2_stat_noisy, p_value_noisy = stats.chisquare(noisy_outputs.numpy().flatten())
# Print the Chi-Square test results for noisy data
print(f"Chi-Square Statistic (Noisy Data): {chi2_stat_noisy}")
print(f"P-Value (Noisy Data): {p_value_noisy}")
# Check if the null hypothesis can be rejected at 1% significance level for noisy data
if p_value_noisy < alpha:
print("Reject the null hypothesis: The noisy data does not follow the standard normal distribution.")
else:
print("Fail to reject the null hypothesis: The noisy data follows the standard normal distribution.")
# Test with increasing noise levels and increasing test set size
noise_levels = [0.1, 0.2, 0.3, 0.4, 0.5]
test_set_sizes = [100, 200, 500, 1000]
results = []
for noise in noise_levels:
for size in test_set_sizes:
noisy_data_test = simulate_sir_noisy(lam, mu, I0, L, noise, noise**2, size, N)[0]
noisy_data_test_tensor = torch.tensor(noisy_data_test, dtype=torch.float32)
with torch.no_grad():
test_outputs = model(noisy_data_test_tensor)
chi2_stat_test, p_value_test = stats.chisquare(test_outputs.numpy().flatten())
results.append({
"noise_level": noise,
"test_set_size": size,
"chi2_stat": chi2_stat_test,
"p_value": p_value_test
})
# Display the results
for result in results:
print(f"Noise Level: {result['noise_level']}, Test Set Size: {result['test_set_size']}, Chi-Square Statistic: {result['chi2_stat']}, P-Value: {result['p_value']}")
Chi-Square Statistic (Noisy Data): 14348.285539387874 P-Value (Noisy Data): 0.0 Reject the null hypothesis: The noisy data does not follow the standard normal distribution. Noise Level: 0.1, Test Set Size: 100, Chi-Square Statistic: 18168.71035071379, P-Value: 0.0 Noise Level: 0.1, Test Set Size: 200, Chi-Square Statistic: 56977.410282305675, P-Value: 0.0 Noise Level: 0.1, Test Set Size: 500, Chi-Square Statistic: 451278.2463351822, P-Value: 0.0 Noise Level: 0.1, Test Set Size: 1000, Chi-Square Statistic: 1513669.8573468588, P-Value: 0.0 Noise Level: 0.2, Test Set Size: 100, Chi-Square Statistic: 10036.911510022237, P-Value: 0.0 Noise Level: 0.2, Test Set Size: 200, Chi-Square Statistic: 52594.09290309042, P-Value: 0.0 Noise Level: 0.2, Test Set Size: 500, Chi-Square Statistic: 310265.118781847, P-Value: 0.0 Noise Level: 0.2, Test Set Size: 1000, Chi-Square Statistic: 957285.3183569491, P-Value: 0.0 Noise Level: 0.3, Test Set Size: 100, Chi-Square Statistic: 23137.732209473706, P-Value: 0.0 Noise Level: 0.3, Test Set Size: 200, Chi-Square Statistic: 34962.803921994986, P-Value: 0.0 Noise Level: 0.3, Test Set Size: 500, Chi-Square Statistic: 220547.6584758056, P-Value: 0.0 Noise Level: 0.3, Test Set Size: 1000, Chi-Square Statistic: 801030.9687531041, P-Value: 0.0 Noise Level: 0.4, Test Set Size: 100, Chi-Square Statistic: 25051.038569875094, P-Value: 0.0 Noise Level: 0.4, Test Set Size: 200, Chi-Square Statistic: 32720.212787365686, P-Value: 0.0 Noise Level: 0.4, Test Set Size: 500, Chi-Square Statistic: 417541.9581649498, P-Value: 0.0 Noise Level: 0.4, Test Set Size: 1000, Chi-Square Statistic: 568316.7527019057, P-Value: 0.0 Noise Level: 0.5, Test Set Size: 100, Chi-Square Statistic: 49801.61596268016, P-Value: 0.0 Noise Level: 0.5, Test Set Size: 200, Chi-Square Statistic: 29737.663633781835, P-Value: 0.0 Noise Level: 0.5, Test Set Size: 500, Chi-Square Statistic: 207126.76490078057, P-Value: 0.0 Noise Level: 0.5, Test Set Size: 1000, Chi-Square Statistic: 616828.6061624703, P-Value: 0.0
Results¶
Test with increasing noise levels and test set sizes, and comment on the findings.
In [29]:
# Test with increasing noise levels and increasing test set sizes
noise_levels = [0.1, 0.2, 0.3, 0.4, 0.5]
test_set_sizes = [100, 200, 500, 1000]
results = []
for noise in noise_levels:
for size in test_set_sizes:
noisy_data_test = simulate_sir_noisy(lam, mu, I0, L, noise, noise**2, size, N)[0]
noisy_data_test_tensor = torch.tensor(noisy_data_test, dtype=torch.float32)
with torch.no_grad():
test_outputs = model(noisy_data_test_tensor)
chi2_stat_test, p_value_test = stats.chisquare(test_outputs.numpy().flatten())
results.append({
"noise_level": noise,
"test_set_size": size,
"chi2_stat": chi2_stat_test,
"p_value": p_value_test
})
# Display the results
for result in results:
print(f"Noise Level: {result['noise_level']}, Test Set Size: {result['test_set_size']}, Chi-Square Statistic: {result['chi2_stat']}, P-Value: {result['p_value']}")
Noise Level: 0.1, Test Set Size: 100, Chi-Square Statistic: 14506.270364824853, P-Value: 0.0 Noise Level: 0.1, Test Set Size: 200, Chi-Square Statistic: 82016.69108116778, P-Value: 0.0 Noise Level: 0.1, Test Set Size: 500, Chi-Square Statistic: 363248.6393957087, P-Value: 0.0 Noise Level: 0.1, Test Set Size: 1000, Chi-Square Statistic: 1337776.0947757147, P-Value: 0.0 Noise Level: 0.2, Test Set Size: 100, Chi-Square Statistic: 10767.540658505433, P-Value: 0.0 Noise Level: 0.2, Test Set Size: 200, Chi-Square Statistic: 40484.933435378625, P-Value: 0.0 Noise Level: 0.2, Test Set Size: 500, Chi-Square Statistic: 279623.3768472299, P-Value: 0.0 Noise Level: 0.2, Test Set Size: 1000, Chi-Square Statistic: 1176851.026197046, P-Value: 0.0 Noise Level: 0.3, Test Set Size: 100, Chi-Square Statistic: 15982.694916609627, P-Value: 0.0 Noise Level: 0.3, Test Set Size: 200, Chi-Square Statistic: 36091.556037124086, P-Value: 0.0 Noise Level: 0.3, Test Set Size: 500, Chi-Square Statistic: 387515.2894235499, P-Value: 0.0 Noise Level: 0.3, Test Set Size: 1000, Chi-Square Statistic: 716550.4344821328, P-Value: 0.0 Noise Level: 0.4, Test Set Size: 100, Chi-Square Statistic: 66768.0326815331, P-Value: 0.0 Noise Level: 0.4, Test Set Size: 200, Chi-Square Statistic: 42658.161048614056, P-Value: 0.0 Noise Level: 0.4, Test Set Size: 500, Chi-Square Statistic: 148172.55713413036, P-Value: 0.0 Noise Level: 0.4, Test Set Size: 1000, Chi-Square Statistic: 538408.0343982365, P-Value: 0.0 Noise Level: 0.5, Test Set Size: 100, Chi-Square Statistic: 22115.53198165923, P-Value: 0.0 Noise Level: 0.5, Test Set Size: 200, Chi-Square Statistic: 55069.70315020295, P-Value: 0.0 Noise Level: 0.5, Test Set Size: 500, Chi-Square Statistic: 304417.2473749254, P-Value: 0.0 Noise Level: 0.5, Test Set Size: 1000, Chi-Square Statistic: 910563.0219575574, P-Value: 0.0
we have demonstrated how to use the MMD loss to detect model mismatch in the context of a simple SIR model.
We have shown that the MMD loss can be used in conjunction with the NLL loss to improve the detection of model mismatch in noisy data.
By comparing the results of the χ2 test on the output of the summary network h(X) with the MMD loss, we have shown that the MMD loss can be used to detect model mismatch in noisy data.
TASK 4 - Sensitivity analysis¶
In [30]:
num_parameters = 3
## Define the Prior Distribution
def prior_distribution(var):
return np.random.normal(0, var, size=(num_samples, num_parameters))
## Vary the Variance of the Prior Distribution
variances = [0.1, 0.5, 1.0]
## Modify Networks to Include Variance as Input
class HyperparameterAwareNetwork(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(HyperparameterAwareNetwork, self).__init__()
self.fc1 = nn.Linear(input_dim + 1, hidden_dim)
self.fc2 = nn.Linear(hidden_dim, output_dim)
def forward(self, x, var):
x = torch.cat((x, var), dim=1)
x = torch.relu(self.fc1(x))
return self.fc2(x)
## Generate Data
def generate_data(num_samples, T, N):
data = []
target = []
for _ in range(num_samples):
lam = np.random.uniform(0.2, 0.9)
mu = np.random.uniform(0.1, 0.7)
I0 = np.random.randint(1, 100)
X, _ = simulate_sir_simple(lam, mu, I0, T, N)
data.append(X.flatten())
target.append([lam, mu, I0])
return np.array(data), np.array(target)
## Create Dataset and DataLoader
class SIRDataset(Dataset):
def __init__(self, X, Y, var):
self.X = torch.tensor(X, dtype=torch.float32)
self.Y = torch.tensor(Y, dtype=torch.float32)
self.var = torch.tensor(var, dtype=torch.float32)
def __len__(self):
return len(self.X)
def __getitem__(self, i):
return self.X[i], self.Y[i], self.var[i]
# Generate data
num_samples = 5000
T = 100
N = 1000
x_train, y_train = generate_data(num_samples, T, N)
var_train = np.random.choice(variances, size=num_samples)
train_dataset = SIRDataset(x_train, y_train, var_train)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
## Train the Model with Varying Variance
input_dim = T * 2
hidden_dim = 128
output_dim = 3
model = HyperparameterAwareNetwork(input_dim, hidden_dim, output_dim)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()
loss_history = []
num_epochs = 50
for epoch in range(num_epochs):
model.train()
epoch_loss = 0.0
for inputs, targets, var in train_loader:
optimizer.zero_grad()
outputs = model(inputs, var.unsqueeze(1))
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
loss_history.append(epoch_loss / len(train_loader))
print(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss / len(train_loader)}")
## Visualize Training Loss
plt.figure(figsize=(8, 4))
plt.plot(loss_history, label='Training Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training Loss History')
plt.legend()
plt.show()
## Evaluate the Model on Test Set
x_test, y_test = generate_data(num_samples=1000, T=T, N=N)
var_test = np.random.choice(variances, size=1000)
test_dataset = SIRDataset(x_test, y_test, var_test)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
model.eval()
preds = []
targets = []
with torch.no_grad():
for inputs, target, var in test_loader:
output = model(inputs, var.unsqueeze(1))
preds.append(output.numpy())
targets.append(target.numpy())
preds = np.vstack(preds)
targets = np.vstack(targets)
## Analyze the Posterior Changes
cols = ['λ', 'μ', 'I_0']
pred_df = pd.DataFrame(preds, columns=cols)
target_df = pd.DataFrame(targets, columns=cols)
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
for i, param in enumerate(cols):
sns.scatterplot(x=target_df[param], y=pred_df[param], alpha=0.7, ax=axes[i])
axes[i].plot([target_df[param].min(), target_df[param].max()],
[target_df[param].min(), target_df[param].max()],
color='red', linestyle='--')
axes[i].set_xlabel(f'True {param}')
axes[i].set_ylabel(f'Predicted {param}')
axes[i].set_title(f'Target vs Predicted for {param}')
plt.tight_layout()
plt.show()
Epoch 1/50, Loss: 307.10729482505894 Epoch 2/50, Loss: 75.16122470324552 Epoch 3/50, Loss: 41.747230167630356 Epoch 4/50, Loss: 30.866556698762917 Epoch 5/50, Loss: 25.877465761160547 Epoch 6/50, Loss: 23.23632894588422 Epoch 7/50, Loss: 20.607949727698216 Epoch 8/50, Loss: 17.988561811326424 Epoch 9/50, Loss: 15.754390746732302 Epoch 10/50, Loss: 13.977813105039958 Epoch 11/50, Loss: 12.651459078245525 Epoch 12/50, Loss: 11.826582968989506 Epoch 13/50, Loss: 10.48136539700665 Epoch 14/50, Loss: 9.88231566586072 Epoch 15/50, Loss: 8.678396315514286 Epoch 16/50, Loss: 7.7886459374729595 Epoch 17/50, Loss: 7.100523661963547 Epoch 18/50, Loss: 6.508655662778057 Epoch 19/50, Loss: 5.898692831208434 Epoch 20/50, Loss: 5.099863007098814 Epoch 21/50, Loss: 4.651367087907429 Epoch 22/50, Loss: 4.367717062370686 Epoch 23/50, Loss: 4.076854877834078 Epoch 24/50, Loss: 3.3162816204602206 Epoch 25/50, Loss: 2.836564531809167 Epoch 26/50, Loss: 2.495497361014161 Epoch 27/50, Loss: 2.3008946482139296 Epoch 28/50, Loss: 1.8804572764076763 Epoch 29/50, Loss: 2.0147008805335322 Epoch 30/50, Loss: 1.6517904453639742 Epoch 31/50, Loss: 1.3022733127014547 Epoch 32/50, Loss: 1.0855428953714008 Epoch 33/50, Loss: 1.1530988269214388 Epoch 34/50, Loss: 0.8606566534012179 Epoch 35/50, Loss: 1.0528062932853457 Epoch 36/50, Loss: 3.742676412757439 Epoch 37/50, Loss: 1.1483961879452573 Epoch 38/50, Loss: 0.56241126188749 Epoch 39/50, Loss: 0.5231969337674636 Epoch 40/50, Loss: 0.4121293133195442 Epoch 41/50, Loss: 0.4305707878327068 Epoch 42/50, Loss: 0.586849986563755 Epoch 43/50, Loss: 0.5934967462775074 Epoch 44/50, Loss: 0.9563818421167664 Epoch 45/50, Loss: 0.3365714142594156 Epoch 46/50, Loss: 0.2704723132184789 Epoch 47/50, Loss: 0.415138538691062 Epoch 48/50, Loss: 0.3461465605452091 Epoch 49/50, Loss: 0.35321289555558677 Epoch 50/50, Loss: 0.39679819062540805
As we can see, when the parameters are varied within the specified ranges, the dynamics of the SIR model can change significantly.
When changing lambda and mu, the change was drastic.
when changing I_0, the change was less.